import json
import math
import os
import os.path as osp
import sys
import time
from datetime import datetime
from time import time
from typing import Iterable, List, Tuple

import cv2
import matplotlib.pyplot as plt
import numpy as np
import sklearn.metrics
import torch
from matplotlib import rcParams
from matplotlib.axes import Axes
from MinkowskiEngine import SparseTensor
from MinkowskiEngine.utils import sparse_quantize
from nuscenes import NuScenes
from nuscenes.lidarseg.lidarseg_utils import (
    colormap_to_colors,
    create_lidarseg_legend,
    get_labels_in_coloring,
    get_stats,
    paint_points_label,
    plt_to_cv2,
)
from nuscenes.utils.color_map import get_colormap
from nuscenes.utils.data_classes import Box, LidarPointCloud, RadarPointCloud
from nuscenes.utils.geometry_utils import (
    BoxVisibility,
    box_in_image,
    transform_matrix,
    view_points,
)
from nuscenes.utils.map_mask import MapMask
from PIL import Image
from pyquaternion import Quaternion
from sklearn import datasets
from sklearn.manifold import TSNE
from tqdm import tqdm

from model import (
    DA,
    Classifier,
    DilationFeatureExtractor,
    DinoVitFeatureExtractor,
    MinkUNet,
    PPKTFeatureExtractor,
    Preprocessing,
    ResNetFCN,
    SegNet,
    VoxelNet,
)
from utils.read_config import generate_config

config = generate_config("config/distil.yaml")

# config["images_encoder"] = "resnet50"
# config["image_weights"] = "moco_v2"
# model_images = DilationFeatureExtractor(config, preprocessing=Preprocessing())

model_points = MinkUNet(1, config["model_n_out"], config)
model_images = SegNet(config, preprocessing=Preprocessing())
model_da = DA(config["model_n_out"], config)

for param in model_points.parameters():
    param.requires_grad = False
for param in model_images.parameters():
    param.requires_grad = False
for param in model_da.parameters():
    param.requires_grad = False

model_points.to(0)
model_images.to(0)
model_da.to(0)

img_checkpoint = torch.load("imgmodel_segnet_83.5.ckpt", map_location="cpu")
img_state = {
    k.replace("model_images.", ""): v
    for k, v in img_checkpoint["state_dict"].items()
    if k.startswith("model_images.")
}
model_images.load_state_dict(img_state, strict=True)

points_checkpoint = torch.load("output/minkunet_slidr_1gpu.pt", map_location="cpu")
pnt_out_path = "slidr.png"

points_checkpoint = torch.load(
    "/root/autodl-tmp/SLidR-2/output/slidr_with_da_3l_relu_with_final/nuscenes/100523-1559/lightning_logs/version_0/checkpoints/lidarmodel-epoch= 49-val_loss= 3.5542.ckpt",
    map_location="cpu",
)
pnt_out_path = "UDAKD.png"

points_checkpoint = torch.load(
    "output/distil_12_MSE_0.5_KL_0.05_da_0.00005_3l_no_final_segnet/100623-1603/lightning_logs/version_0/checkpoints/lidarmodel-epoch=29-m_IoU=0.6372.ckpt",
    map_location="cpu",
)
pnt_out_path = "FSKD.png"

nusc = NuScenes(version="v1.0-mini", dataroot="datasets/nuscenes", verbose=True)
nusc.nusc = nusc
# my_sample = nusc.sample[87]
# my_sample = nusc.sample[99]
# my_sample = nusc.sample[11]
# my_sample = nusc.sample[17]
# my_sample = nusc.sample[199]
# my_sample = nusc.sample[295]
my_sample = nusc.sample[399]


def load_state_with_same_shape(model, weights):
    """
    Load common weights in two similar models
    (for instance between a pretraining and a downstream training)
    """
    model_state = model.state_dict()
    if list(weights.keys())[0].startswith("model."):
        weights = {k.partition("model.")[2]: weights[k] for k in weights.keys()}

    if list(weights.keys())[0].startswith("model_points."):
        weights = {k.partition("model_points.")[2]: weights[k] for k in weights.keys()}

    if list(weights.keys())[0].startswith("module."):
        print("Loading multigpu weights with module. prefix...")
        weights = {k.partition("module.")[2]: weights[k] for k in weights.keys()}

    if list(weights.keys())[0].startswith("encoder."):
        print("Loading multigpu weights with encoder. prefix...")
        weights = {k.partition("encoder.")[2]: weights[k] for k in weights.keys()}

    filtered_weights = {
        k: v
        for k, v in weights.items()
        if (k in model_state and v.size() == model_state[k].size())
    }
    removed_weights = {
        k: v
        for k, v in weights.items()
        if not (k in model_state and v.size() == model_state[k].size())
    }
    print("Loading weights:" + ", ".join(filtered_weights.keys()))
    print("")
    print("Not loading weights:" + ", ".join(removed_weights.keys()))
    return filtered_weights


print("Training with pretrained model")
if "config" in points_checkpoint:
    for cfg in ("voxel_size", "cylindrical_coordinates"):
        assert points_checkpoint["config"][cfg] == config[cfg], (
            f"{cfg} is not consistant. "
            f"Checkpoint: {points_checkpoint['config'][cfg]}, "
            f"Config: {config[cfg]}."
        )
if set(points_checkpoint.keys()) == set(
    ["epoch", "model", "optimizer", "train_criterion"]
):
    print("Pre-trained weights are coming from DepthContrast.")
    pretraining_epochs = points_checkpoint["epoch"]
    print(f"==> Number of pre-training epochs {pretraining_epochs}")
    points_checkpoint = points_checkpoint["model"]
    if list(points_checkpoint.keys())[0].startswith("module."):
        print("Loading multigpu weights with module. prefix...")
        points_checkpoint = {
            k.partition("module.")[2]: points_checkpoint[k]
            for k in points_checkpoint.keys()
        }
    voxel_net_suffix = "trunk.2."
    points_checkpoint = {
        key.partition(voxel_net_suffix)[2]: points_checkpoint[key]
        for key in points_checkpoint.keys()
        if key.startswith(voxel_net_suffix)
    }
    print(f"==> Number of loaded weight blobs {len(points_checkpoint)}")
    points_checkpoint = {"model_points": points_checkpoint}
key = "model_points" if "model_points" in points_checkpoint else "state_dict"
filtered_weights = load_state_with_same_shape(model_points, points_checkpoint[key])
model_dict = model_points.state_dict()
model_dict.update(filtered_weights)
model_points.load_state_dict(model_dict)

# points_state = {k.replace("model_points.", ""): v for k,
#                 v in points_checkpoint["state_dict"].items() if k.startswith("model_points.")}
# model_points.load_state_dict(points_state, strict=True)

# da_state = {k.replace("model_da.", ""): v for k,
#             v in points_checkpoint["state_dict"].items() if k.startswith("model_da.")}
# model_da.load_state_dict(da_state, strict=True)


sample_data_token = my_sample["data"]["LIDAR_TOP"]

# nusc.render_sample_data(sample_data_token,
#                         with_anns=False,
#                         show_lidarseg=True, out_path="tmp.png")

# tmp = nusc.get("sample_data", 'd219ffa9b4ce492c8b8059db9d8b3166')
# print(tmp["filename"])

sample_data_token = sample_data_token
with_anns = False
box_vis_level = BoxVisibility.ANY
axes_limit = 100
ax = None
nsweeps = 1
out_path = "tmp.png"
underlay_map = False
use_flat_vehicle_coordinates = True
show_lidarseg = False
show_lidarseg_legend = False
filter_lidarseg_labels: List = None
lidarseg_preds_bin_path = None
verbose = True

# tmp_img = cv2.imread(
#     "./datasets/nuscenes/samples/CAM_FRONT/n008-2018-08-28-16-43-51-0400__CAM_FRONT__1535489300012404.jpg")
# # tmp_img = cv2.imread(
# #     "./datasets/nuscenes/samples/CAM_BACK/n008-2018-08-28-16-43-51-0400__CAM_BACK__1535489300037558.jpg")
# cv2.imwrite("tmp_img.png", tmp_img)
# tmp_img = cv2.resize(tmp_img, (1600, 896))
# tmp_img = torch.tensor(tmp_img, dtype=torch.float32).permute(2, 0, 1)/255
# tmp_img = tmp_img.unsqueeze(0)

# tmp_feat = model_images(tmp_img.to(0)).squeeze(0)
# tmp_feat = torch.linalg.vector_norm(tmp_feat, dim=0)
# tmp_feat = (tmp_feat-tmp_feat.min()) / \
#     (tmp_feat.max()-tmp_feat.min())*255
# tmp_feat = np.array(tmp_feat.cpu(), dtype=np.uint8)
# tmp_feat = cv2.applyColorMap(tmp_feat, cv2.COLORMAP_JET)
# tmp_feat = cv2.resize(tmp_feat, (1600, 900))
# cv2.imwrite("tmp_feat.png", tmp_feat)

# exit(0)

points = LidarPointCloud.from_file(
    nusc.dataroot + "/" + nusc.get("sample_data", sample_data_token)["filename"]
).points.T
# get the points (4th coordinate is the point intensity)
pc = points[:, :3]
# if self.labels:
#     lidarseg_labels_filename = os.path.join(
#         self.nusc.dataroot, self.nusc.get(
#             "lidarseg", lidar_token)["filename"])
#     points_labels = np.fromfile(
#         lidarseg_labels_filename, dtype=np.uint8)
pc = torch.tensor(pc)
# apply the transforms (augmentation)

voxel_size = config["voxel_size"]

# Transform to cylinder coordinate and scale for given voxel size
x, y, z = pc.T
rho = torch.sqrt(x**2 + y**2) / voxel_size
# corresponds to a split each 1°
phi = torch.atan2(y, x) * 180 / np.pi
z = z / voxel_size
coords_aug = torch.cat((rho[:, None], phi[:, None], z[:, None]), 1)


# Voxelization
discrete_coords, indexes, inverse_indexes = sparse_quantize(
    coords_aug, return_index=True, return_inverse=True
)
# use those voxels features
unique_feats = torch.tensor(points[indexes][:, 3:])

discrete_coords = torch.nn.functional.pad(
    input=discrete_coords, pad=(1, 0, 0, 0), mode="constant", value=0
).to(0)
unique_feats = unique_feats.to(0)

sparse_input = SparseTensor(unique_feats, discrete_coords)
output_points = model_points(sparse_input)


def map_pointcloud_to_image(
    self,
    pointsensor_token: str,
    camera_token: str,
    min_dist: float = 1.0,
    render_intensity: bool = False,
    show_lidarseg: bool = False,
    filter_lidarseg_labels: List = None,
    lidarseg_preds_bin_path: str = None,
) -> Tuple:
    """
    Given a point sensor (lidar/radar) token and camera sample_data token, load pointcloud and map it to the image
    plane.
    :param pointsensor_token: Lidar/radar sample_data token.
    :param camera_token: Camera sample_data token.
    :param min_dist: Distance from the camera below which points are discarded.
    :param render_intensity: Whether to render lidar intensity instead of point depth.
    :param show_lidarseg: Whether to render lidar intensity instead of point depth.
    :param filter_lidarseg_labels: Only show lidar points which belong to the given list of classes. If None
        or the list is empty, all classes will be displayed.
    :param lidarseg_preds_bin_path: A path to the .bin file which contains the user's lidar segmentation
                                    predictions for the sample.
    :return (pointcloud <np.float: 2, n)>, coloring <np.float: n>, image <Image>).
    """
    cam = self.nusc.get("sample_data", camera_token)
    pointsensor = self.nusc.get("sample_data", pointsensor_token)
    pcl_path = osp.join(self.nusc.dataroot, pointsensor["filename"])
    if pointsensor["sensor_modality"] == "lidar":
        if show_lidarseg:
            assert hasattr(
                self.nusc, "lidarseg"
            ), "Error: nuScenes-lidarseg not installed!"
            # Ensure that lidar pointcloud is from a keyframe.
            assert pointsensor[
                "is_key_frame"
            ], "Error: Only pointclouds which are keyframes have lidar segmentation labels. Rendering aborted."
            assert not render_intensity, (
                "Error: Invalid options selected. You can only select either "
                "render_intensity or show_lidarseg, not both."
            )
        pc = LidarPointCloud.from_file(pcl_path)
    else:
        pc = RadarPointCloud.from_file(pcl_path)
    im = Image.open(osp.join(self.nusc.dataroot, cam["filename"]))
    # Points live in the point sensor frame. So they need to be transformed via global to the image plane.
    # First step: transform the pointcloud to the ego vehicle frame for the timestamp of the sweep.
    cs_record = self.nusc.get(
        "calibrated_sensor", pointsensor["calibrated_sensor_token"]
    )
    pc.rotate(Quaternion(cs_record["rotation"]).rotation_matrix)
    pc.translate(np.array(cs_record["translation"]))
    # Second step: transform from ego to the global frame.
    poserecord = self.nusc.get("ego_pose", pointsensor["ego_pose_token"])
    pc.rotate(Quaternion(poserecord["rotation"]).rotation_matrix)
    pc.translate(np.array(poserecord["translation"]))
    # Third step: transform from global into the ego vehicle frame for the timestamp of the image.
    poserecord = self.nusc.get("ego_pose", cam["ego_pose_token"])
    pc.translate(-np.array(poserecord["translation"]))
    pc.rotate(Quaternion(poserecord["rotation"]).rotation_matrix.T)
    # Fourth step: transform from ego into the camera.
    cs_record = self.nusc.get("calibrated_sensor", cam["calibrated_sensor_token"])
    pc.translate(-np.array(cs_record["translation"]))
    pc.rotate(Quaternion(cs_record["rotation"]).rotation_matrix.T)
    # Fifth step: actually take a "picture" of the point cloud.
    # Grab the depths (camera frame z axis points away from the camera).
    depths = pc.points[2, :]
    if render_intensity:
        assert pointsensor["sensor_modality"] == "lidar", (
            "Error: Can only render intensity for lidar, "
            "not %s!" % pointsensor["sensor_modality"]
        )
        # Retrieve the color from the intensities.
        # Performs arbitary scaling to achieve more visually pleasing results.
        intensities = pc.points[3, :]

        intensities = output_points.F[inverse_indexes]
        intensities = torch.linalg.vector_norm(intensities, dim=1).cpu()
        intensities = np.array(intensities)

        # intensities = (intensities - np.min(intensities)) / \
        #     (np.max(intensities) - np.min(intensities))
        # intensities = intensities ** 0.1
        # intensities = np.maximum(0, intensities - 0.5)
        coloring = intensities
    elif show_lidarseg:
        assert pointsensor["sensor_modality"] == "lidar", (
            "Error: Can only render lidarseg labels for lidar, "
            "not %s!" % pointsensor["sensor_modality"]
        )
        if lidarseg_preds_bin_path:
            sample_token = self.nusc.get("sample_data", pointsensor_token)[
                "sample_token"
            ]
            lidarseg_labels_filename = lidarseg_preds_bin_path
            assert os.path.exists(lidarseg_labels_filename), (
                "Error: Unable to find {} to load the predictions for sample token {} (lidar "
                "sample data token {}) from.".format(
                    lidarseg_labels_filename, sample_token, pointsensor_token
                )
            )
        else:
            # Ensure lidarseg.json is not empty (e.g. in case of v1.0-test).
            if len(self.nusc.lidarseg) > 0:
                lidarseg_labels_filename = osp.join(
                    self.nusc.dataroot,
                    self.nusc.get("lidarseg", pointsensor_token)["filename"],
                )
            else:
                lidarseg_labels_filename = None
        if lidarseg_labels_filename:
            # Paint each label in the pointcloud with a RGBA value.
            coloring = paint_points_label(
                lidarseg_labels_filename,
                filter_lidarseg_labels,
                self.nusc.lidarseg_name2idx_mapping,
                self.nusc.colormap,
            )
        else:
            coloring = depths
            print(
                "Warning: There are no lidarseg labels in {}. Points will be colored according to distance "
                "from the ego vehicle instead.".format(self.nusc.version)
            )
    else:
        # Retrieve the color from the depth.
        coloring = depths
    # Take the actual picture (matrix multiplication with camera-matrix + renormalization).
    points = view_points(
        pc.points[:3, :], np.array(cs_record["camera_intrinsic"]), normalize=True
    )
    # Remove points that are either outside or behind the camera. Leave a margin of 1 pixel for aesthetic reasons.
    # Also make sure points are at least 1m in front of the camera to avoid seeing the lidar points on the camera
    # casing for non-keyframes which are slightly out of sync.
    mask = np.ones(depths.shape[0], dtype=bool)
    mask = np.logical_and(mask, depths > min_dist)
    mask = np.logical_and(mask, points[0, :] > 1)
    mask = np.logical_and(mask, points[0, :] < im.size[0] - 1)
    mask = np.logical_and(mask, points[1, :] > 1)
    mask = np.logical_and(mask, points[1, :] < im.size[1] - 1)
    points = points[:, mask]
    coloring = coloring[mask]
    return points, coloring, im


def render_pointcloud_in_image(
    self,
    sample_token: str,
    dot_size: int = 5,
    pointsensor_channel: str = "LIDAR_TOP",
    camera_channel: str = "CAM_FRONT",
    out_path: str = None,
    render_intensity: bool = False,
    show_lidarseg: bool = False,
    filter_lidarseg_labels: List = None,
    ax: Axes = None,
    show_lidarseg_legend: bool = False,
    verbose: bool = True,
    lidarseg_preds_bin_path: str = None,
):
    """
    Scatter-plots a pointcloud on top of image.
    :param sample_token: Sample token.
    :param dot_size: Scatter plot dot size.
    :param pointsensor_channel: RADAR or LIDAR channel name, e.g. 'LIDAR_TOP'.
    :param camera_channel: Camera channel name, e.g. 'CAM_FRONT'.
    :param out_path: Optional path to save the rendered figure to disk.
    :param render_intensity: Whether to render lidar intensity instead of point depth.
    :param show_lidarseg: Whether to render lidarseg labels instead of point depth.
    :param filter_lidarseg_labels: Only show lidar points which belong to the given list of classes.
    :param ax: Axes onto which to render.
    :param show_lidarseg_legend: Whether to display the legend for the lidarseg labels in the frame.
    :param verbose: Whether to display the image in a window.
    :param lidarseg_preds_bin_path: A path to the .bin file which contains the user's lidar segmentation
                                    predictions for the sample.
    """
    sample_record = self.nusc.get("sample", sample_token)
    # Here we just grab the front camera and the point sensor.
    pointsensor_token = sample_record["data"][pointsensor_channel]
    camera_token = sample_record["data"][camera_channel]
    points, coloring, im = map_pointcloud_to_image(
        self,
        pointsensor_token,
        camera_token,
        render_intensity=render_intensity,
        show_lidarseg=show_lidarseg,
        filter_lidarseg_labels=filter_lidarseg_labels,
        lidarseg_preds_bin_path=lidarseg_preds_bin_path,
    )
    # Init axes.
    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=(9, 16))
        if lidarseg_preds_bin_path:
            fig.canvas.manager.set_window_title(sample_token + "(predictions)")
        else:
            fig.canvas.manager.set_window_title(sample_token)
    else:  # Set title on if rendering as part of render_sample.
        ax.set_title(camera_channel)
    ax.imshow(im)

    coloring = (coloring - coloring.min()) / (coloring.max() - coloring.min()) * 255
    coloring = np.array(coloring, dtype=np.uint8)
    coloring = cv2.applyColorMap(coloring, cv2.COLORMAP_JET)
    coloring = coloring.squeeze()
    coloring = coloring[:, (2, 1, 0)]
    coloring = np.divide(coloring, 255, dtype=np.float32)

    ax.scatter(points[0, :], points[1, :], c=coloring, s=dot_size)
    ax.axis("off")
    # Produce a legend with the unique colors from the scatter.
    if pointsensor_channel == "LIDAR_TOP" and show_lidarseg and show_lidarseg_legend:
        # Since the labels are stored as class indices, we get the RGB colors from the colormap in an array where
        # the position of the RGB color corresponds to the index of the class it represents.
        color_legend = colormap_to_colors(
            self.nusc.colormap, self.nusc.lidarseg_name2idx_mapping
        )
        # If user does not specify a filter, then set the filter to contain the classes present in the pointcloud
        # after it has been projected onto the image; this will allow displaying the legend only for classes which
        # are present in the image (instead of all the classes).
        if filter_lidarseg_labels is None:
            filter_lidarseg_labels = get_labels_in_coloring(color_legend, coloring)
        create_lidarseg_legend(
            filter_lidarseg_labels,
            self.nusc.lidarseg_idx2name_mapping,
            self.nusc.colormap,
        )
    if out_path is not None:
        plt.savefig(out_path, bbox_inches="tight", pad_inches=0, dpi=960)
    if verbose:
        plt.show()


render_pointcloud_in_image(
    nusc,
    my_sample["token"],
    pointsensor_channel="LIDAR_TOP",
    camera_channel="CAM_FRONT",
    render_intensity=True,
    show_lidarseg=False,
    # filter_lidarseg_labels=[22, 23, 24],
    show_lidarseg_legend=False,
    out_path=pnt_out_path,
)

exit(0)


# def render_ego_centric_map(
#         sample_data_token: str,
#         axes_limit: float = 40,
#         ax: Axes = None) -> None:
#     """
#     Render map centered around the associated ego pose.
#     :param sample_data_token: Sample_data token.
#     :param axes_limit: Axes limit measured in meters.
#     :param ax: Axes onto which to render.
#     """

#     def crop_image(image: np.array,
#                    x_px: int,
#                    y_px: int,
#                    axes_limit_px: int) -> np.array:
#         x_min = int(x_px - axes_limit_px)
#         x_max = int(x_px + axes_limit_px)
#         y_min = int(y_px - axes_limit_px)
#         y_max = int(y_px + axes_limit_px)

#         cropped_image = image[y_min:y_max, x_min:x_max]

#         return cropped_image

#     # Get data.
#     sd_record = nusc.get('sample_data', sample_data_token)
#     sample = nusc.get('sample', sd_record['sample_token'])
#     scene = nusc.get('scene', sample['scene_token'])
#     log = nusc.get('log', scene['log_token'])
#     map_ = nusc.get('map', log['map_token'])
#     map_mask = map_['mask']
#     pose = nusc.get('ego_pose', sd_record['ego_pose_token'])

#     # Retrieve and crop mask.
#     pixel_coords = map_mask.to_pixel_coords(
#         pose['translation'][0], pose['translation'][1])
#     scaled_limit_px = int(axes_limit * (1.0 / map_mask.resolution))
#     mask_raster = map_mask.mask()
#     cropped = crop_image(mask_raster, pixel_coords[0], pixel_coords[1], int(
#         scaled_limit_px * math.sqrt(2)))

#     # Rotate image.
#     ypr_rad = Quaternion(pose['rotation']).yaw_pitch_roll
#     yaw_deg = -math.degrees(ypr_rad[0])
#     rotated_cropped = np.array(Image.fromarray(cropped).rotate(yaw_deg))

#     # Crop image.
#     ego_centric_map = crop_image(rotated_cropped, rotated_cropped.shape[1] / 2,
#                                  rotated_cropped.shape[0] / 2,
#                                  scaled_limit_px)

#     # Init axes and show image.
#     # Set background to white and foreground (semantic prior) to gray.
#     if ax is None:
#         _, ax = plt.subplots(1, 1, figsize=(9, 9))
#     ego_centric_map[ego_centric_map == map_mask.foreground] = 125
#     ego_centric_map[ego_centric_map == map_mask.background] = 255
#     ax.imshow(ego_centric_map, extent=[-axes_limit, axes_limit, -axes_limit, axes_limit],
#               cmap='gray', vmin=0, vmax=255)


# sd_record = nusc.get('sample_data', sample_data_token)
# sensor_modality = sd_record['sensor_modality']
# if sensor_modality in ['lidar', 'radar']:
#     sample_rec = nusc.get('sample', sd_record['sample_token'])
#     chan = sd_record['channel']
#     ref_chan = 'LIDAR_TOP'
#     ref_sd_token = sample_rec['data'][ref_chan]
#     ref_sd_record = nusc.get('sample_data', ref_sd_token)
#     if sensor_modality == 'lidar':
#         if show_lidarseg:
#             assert hasattr(
#                 nusc, 'lidarseg'), 'Error: nuScenes-lidarseg not installed!'
#             # Ensure that lidar pointcloud is from a keyframe.
#             assert sd_record['is_key_frame'], \
#                 'Error: Only pointclouds which are keyframes have lidar segmentation labels. Rendering aborted.'
#             assert nsweeps == 1, \
#                 'Error: Only pointclouds which are keyframes have lidar segmentation labels; nsweeps should ' \
#                 'be set to 1.'
#             # Load a single lidar point cloud.
#             pcl_path = osp.join(nusc.dataroot, ref_sd_record['filename'])
#             pc = LidarPointCloud.from_file(pcl_path)
#         else:
#             # Get aggregated lidar point cloud in lidar frame.
#             pc, times = LidarPointCloud.from_file_multisweep(nusc, sample_rec, chan, ref_chan,
#                                                              nsweeps=nsweeps)
#         velocities = None
#     else:
#         # Get aggregated radar point cloud in reference frame.
#         # The point cloud is transformed to the reference frame for visualization purposes.
#         pc, times = RadarPointCloud.from_file_multisweep(
#             nusc, sample_rec, chan, ref_chan, nsweeps=nsweeps)
#         # Transform radar velocities (x is front, y is left), as these are not transformed when loading the
#         # point cloud.
#         radar_cs_record = nusc.get(
#             'calibrated_sensor', sd_record['calibrated_sensor_token'])
#         ref_cs_record = nusc.get(
#             'calibrated_sensor', ref_sd_record['calibrated_sensor_token'])
#         velocities = pc.points[8:10, :]  # Compensated velocity
#         velocities = np.vstack((velocities, np.zeros(pc.points.shape[1])))
#         velocities = np.dot(Quaternion(
#             radar_cs_record['rotation']).rotation_matrix, velocities)
#         velocities = np.dot(Quaternion(
#             ref_cs_record['rotation']).rotation_matrix.T, velocities)
#         velocities[2, :] = np.zeros(pc.points.shape[1])
#     # By default we render the sample_data top down in the sensor frame.
#     # This is slightly inaccurate when rendering the map as the sensor frame may not be perfectly upright.
#     # Using use_flat_vehicle_coordinates we can render the map in the ego frame instead.
#     if use_flat_vehicle_coordinates:
#         # Retrieve transformation matrices for reference point cloud.
#         cs_record = nusc.get(
#             'calibrated_sensor', ref_sd_record['calibrated_sensor_token'])
#         pose_record = nusc.get(
#             'ego_pose', ref_sd_record['ego_pose_token'])
#         ref_to_ego = transform_matrix(translation=cs_record['translation'],
#                                       rotation=Quaternion(cs_record["rotation"]))
#         # Compute rotation between 3D vehicle pose and "flat" vehicle pose (parallel to global z plane).
#         ego_yaw = Quaternion(pose_record['rotation']).yaw_pitch_roll[0]
#         rotation_vehicle_flat_from_vehicle = np.dot(
#             Quaternion(scalar=np.cos(ego_yaw / 2),
#                        vector=[0, 0, np.sin(ego_yaw / 2)]).rotation_matrix,
#             Quaternion(pose_record['rotation']).inverse.rotation_matrix)
#         vehicle_flat_from_vehicle = np.eye(4)
#         vehicle_flat_from_vehicle[:3, :3] = rotation_vehicle_flat_from_vehicle
#         viewpoint = np.dot(vehicle_flat_from_vehicle, ref_to_ego)

#         tmp_qua = [0.2406501, 0.4813002, 0.3609751, 0.7616657]
#         qua = [tmp_qua[3]]+tmp_qua[0:3]

#         viewpoint = np.dot(viewpoint, transform_matrix(
#             rotation=Quaternion(qua)))
#     else:
#         viewpoint = np.eye(4)
#     # Init axes.
#     if ax is None:
#         _, ax = plt.subplots(1, 1, figsize=(9, 9))
#     # Render map if requested.
#     if underlay_map:
#         assert use_flat_vehicle_coordinates, 'Error: underlay_map requires use_flat_vehicle_coordinates, as ' \
#                                              'otherwise the location does not correspond to the map!'
#         render_ego_centric_map(
#             sample_data_token=sample_data_token, axes_limit=axes_limit, ax=ax)
#     # Show point cloud.
#     points = view_points(pc.points[:3, :], viewpoint, normalize=False)
#     dists = np.sqrt(np.sum(pc.points[:2, :] ** 2, axis=0))

#     # intensities = pc.points[3, :]
#     # intensities = (intensities - np.min(intensities)) / \
#     #     (np.max(intensities) - np.min(intensities))
#     # intensities = intensities ** 0.1
#     # intensities = np.maximum(0, intensities - 0.5)
#     # dists = intensities
#     dists = np.log(pc.points[3, :]+1)

#     if sensor_modality == 'lidar' and show_lidarseg:
#         # Load labels for pointcloud.
#         if lidarseg_preds_bin_path:
#             sample_token = nusc.get('sample_data', sample_data_token)[
#                 'sample_token']
#             lidarseg_labels_filename = lidarseg_preds_bin_path
#             assert os.path.exists(lidarseg_labels_filename), \
#                 'Error: Unable to find {} to load the predictions for sample token {} (lidar ' \
#                 'sample data token {}) from.'.format(
#                     lidarseg_labels_filename, sample_token, sample_data_token)
#         else:
#             # Ensure lidarseg.json is not empty (e.g. in case of v1.0-test).
#             if len(nusc.lidarseg) > 0:
#                 lidarseg_labels_filename = osp.join(nusc.dataroot,
#                                                     nusc.get('lidarseg', sample_data_token)['filename'])
#             else:
#                 lidarseg_labels_filename = None
#         if lidarseg_labels_filename:
#             # Paint each label in the pointcloud with a RGBA value.
#             colors = paint_points_label(lidarseg_labels_filename, filter_lidarseg_labels,
#                                         nusc.lidarseg_name2idx_mapping, nusc.colormap)
#             if show_lidarseg_legend:
#                 # Since the labels are stored as class indices, we get the RGB colors from the colormap
#                 # in an array where the position of the RGB color corresponds to the index of the class
#                 # it represents.
#                 color_legend = colormap_to_colors(
#                     nusc.colormap, nusc.lidarseg_name2idx_mapping)
#                 # If user does not specify a filter, then set the filter to contain the classes present in
#                 # the pointcloud after it has been projected onto the image; this will allow displaying the
#                 # legend only for classes which are present in the image (instead of all the classes).
#                 if filter_lidarseg_labels is None:
#                     filter_lidarseg_labels = get_labels_in_coloring(
#                         color_legend, colors)
#                 create_lidarseg_legend(filter_lidarseg_labels,
#                                        nusc.lidarseg_idx2name_mapping, nusc.colormap,
#                                        loc='upper left', ncol=1, bbox_to_anchor=(1.05, 1.0))
#         else:
#             colors = np.minimum(1, dists / axes_limit / np.sqrt(2))
#             print('Warning: There are no lidarseg labels in {}. Points will be colored according to distance '
#                   'from the ego vehicle instead.'.format(nusc.version))
#     else:
#         colors = np.minimum(1, dists / axes_limit / np.sqrt(2))
#     point_scale = 0.2 if sensor_modality == 'lidar' else 3.0
#     scatter = ax.scatter(points[0, :], points[1, :], c=colors, s=point_scale)
#     # Show velocities.
#     if sensor_modality == 'radar':
#         points_vel = view_points(
#             pc.points[:3, :] + velocities, viewpoint, normalize=False)
#         deltas_vel = points_vel - points
#         deltas_vel = 6 * deltas_vel  # Arbitrary scaling
#         max_delta = 20
#         deltas_vel = np.clip(deltas_vel, -max_delta,
#                              max_delta)  # Arbitrary clipping
#         colors_rgba = scatter.to_rgba(colors)
#         for i in range(points.shape[1]):
#             ax.arrow(points[0, i], points[1, i], deltas_vel[0, i],
#                      deltas_vel[1, i], color=colors_rgba[i])
#     # # Show ego vehicle.
#     # ax.plot(0, 0, 'x', color='red')
#     # Get boxes in lidar frame.
#     _, boxes, _ = nusc.get_sample_data(ref_sd_token, box_vis_level=box_vis_level,
#                                        use_flat_vehicle_coordinates=use_flat_vehicle_coordinates)
#     # Show boxes.
#     if with_anns:
#         for box in boxes:
#             c = np.array(nusc.get_color(box.name)) / 255.0
#             box.render(ax, view=np.eye(4), colors=(c, c, c))
#     # Limit visible range.
#     ax.set_xlim(-axes_limit, axes_limit)
#     ax.set_ylim(-axes_limit, axes_limit)
# elif sensor_modality == 'camera':
#     # Load boxes and image.
#     data_path, boxes, camera_intrinsic = nusc.get_sample_data(sample_data_token,
#                                                               box_vis_level=box_vis_level)
#     data = Image.open(data_path)
#     # Init axes.
#     if ax is None:
#         _, ax = plt.subplots(1, 1, figsize=(9, 16))
#     # Show image.
#     ax.imshow(data)
#     # Show boxes.
#     if with_anns:
#         for box in boxes:
#             c = np.array(nusc.get_color(box.name)) / 255.0
#             box.render(ax, view=camera_intrinsic,
#                        normalize=True, colors=(c, c, c))
#     # Limit visible range.
#     ax.set_xlim(0, data.size[0])
#     ax.set_ylim(data.size[1], 0)
# else:
#     raise ValueError("Error: Unknown sensor modality!")
# ax.axis('off')
# ax.set_title('{} {labels_type}'.format(
#     sd_record['channel'], labels_type='(predictions)' if lidarseg_preds_bin_path else ''))
# ax.set_aspect('equal')
# if out_path is not None:
#     plt.savefig(out_path, bbox_inches='tight', pad_inches=0, dpi=2000)
# if verbose:
#     plt.show()
